import tensorflow as tf  # 引入TensorFlow库,用于构建和训练神经网络模型
from tensorflow.keras.datasets import mnist  # 引入MNIST手写数字数据集
from tensorflow.keras.layers import Flatten, Dense  # 引入神经网络层:Flatten和Dense
from tensorflow.keras.models import Sequential  # 引入Sequential模型
from tensorflow.keras.utils import to_categorical  # 引入to_categorical函数,用于将标签转换为one-hot编码
import numpy as np  # 引入NumPy库,用于数组操作和数值计算
import gradio as gr
import cv2


# 1. 引入数据,对数据进行预处理
(X_train, y_train), (X_test, y_test) = mnist.load_data()  # 加载MNIST数据集,并将其拆分为训练集和测试集
X_train = X_train / 255.0  # 将训练集像素值归一化到[0,1]范围内,提高模型训练效果
X_test = X_test / 255.0  # 将测试集像素值归一化到[0,1]范围内,提高模型推理效果
y_train = to_categorical(y_train)  # 将训练集标签转换为one-hot编码,便于模型训练
y_test = to_categorical(y_test)  # 将测试集标签转换为one-hot编码,便于模型评估

# 2. 构建网络模型
model = Sequential()  # 创建Sequential模型
model.add(Flatten(input_shape=(28, 28)))  # 添加Flatten层,将二维图像数据展平为一维向量
model.add(Dense(20, activation=tf.nn.relu))  # 添加全连接层,使用ReLU激活函数,该层有20个神经元
model.add(Dense(10, activation=tf.nn.softmax))  # 添加输出层,使用Softmax激活函数,该层有10个神经元,对应10个类别

# 3. 编译模型,指定优化器、损失函数和评估指标
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 4. 训练模型
model.fit(X_train, y_train, epochs=20, validation_split=0.2)  # 训练模型,设置训练轮数为20,验证集比例为20%

# 5. 评估模型
test_loss, test_acc = model.evaluate(X_test, y_test)  # 在测试集上评估模型性能
print('Test accuracy:', test_acc)  # 输出测试集准确率

# 6. 模型推理
X_infer = np.array([X_test[0], X_test[1]])  # 从测试集中选择前两个样本作为推理数据
y_infer = model.predict(X_infer)  # 使用训练好的模型对推理数据进行预测

# 保存模型到硬盘
model_path = './mnist_model'  # 指定保存模型的路径和文件名
tf.keras.models.save_model(model, model_path) # 使用save_model函数将训练好的模型保存到指定路径,实现模型持久化

# 从硬盘加载保存的模型
loaded_model = tf.keras.models.load_model(model_path)  # 使用 load_model 函数从指定路径加载保存的模型

def predict_digit(image):
    image = cv2.resize(image, (28, 28))
    image = image.astype('float32') / 255.0
    image = np.expand_dims(image, axis=0)  # 添加额外的维度，使其成为形状为 (1, 28, 28) 的数据
    predicted_digit = np.argmax(loaded_model.predict(image))
    return str(predicted_digit)

# 创建 Gradio 接口
interface = gr.Interface(
    fn=predict_digit,
    inputs='sketchpad',
    outputs='text',
    title='手写数字识别',
    description='使用 TensorFlow 的 MLP 进行手写数字识别'
)

# 启动 Gradio 接口
interface.launch()
